#include "crypto.h"
#include <stdlib.h>
#include <string.h>

// 大整数结构
typedef struct {
    uint32_t *data;
    size_t size;
    int sign;  // 1 for positive, -1 for negative
} bigint_t;

// 辅助函数
static uint32_t mod_pow(uint32_t base, uint32_t exponent, uint32_t modulus) {
    uint32_t result = 1;
    base = base % modulus;
    
    while (exponent > 0) {
        if (exponent & 1) {
            result = (result * base) % modulus;
        }
        base = (base * base) % modulus;
        exponent >>= 1;
    }
    return result;
}

static int is_prime(uint32_t n) {
    if (n < 2) return 0;
    if (n == 2) return 1;
    if (n % 2 == 0) return 0;
    
    for (uint32_t i = 3; i * i <= n; i += 2) {
        if (n % i == 0) return 0;
    }
    return 1;
}

static uint32_t generate_prime(uint32_t min, uint32_t max) {
    uint32_t prime;
    do {
        prime = (rand() % (max - min + 1)) + min;
        if (prime % 2 == 0) prime++;
    } while (!is_prime(prime));
    return prime;
}

static uint32_t gcd(uint32_t a, uint32_t b) {
    while (b != 0) {
        uint32_t temp = b;
        b = a % b;
        a = temp;
    }
    return a;
}

static uint32_t extended_gcd(uint32_t a, uint32_t b, int *x, int *y) {
    if (b == 0) {
        *x = 1;
        *y = 0;
        return a;
    }
    
    int x1, y1;
    uint32_t result = extended_gcd(b, a % b, &x1, &y1);
    
    *x = y1;
    *y = x1 - (a / b) * y1;
    
    return result;
}

static uint32_t mod_inverse(uint32_t a, uint32_t m) {
    int x = 0, y = 0;
    uint32_t g = extended_gcd(a, m, &x, &y);
    if (g != 1) return 0; // 不存在逆元
    
    return (x % m + m) % m;
}

// 简化的RSA实现（使用32位整数）
int rsa_generate_keypair(rsa_keypair_t *keypair, size_t key_size) {
    if (!keypair) return CRYPTO_ERROR_INVALID_PARAM;
    
    // 生成两个大素数
    uint32_t p = generate_prime(1000, 10000);
    uint32_t q = generate_prime(1000, 10000);
    
    uint32_t n = p * q;
    uint32_t phi = (p - 1) * (q - 1);
    
    // 选择公钥指数e
    uint32_t e = 65537; // 常用的公钥指数
    while (gcd(e, phi) != 1) {
        e += 2;
    }
    
    // 计算私钥指数d
    uint32_t d = mod_inverse(e, phi);
    
    // 分配内存并存储密钥
    keypair->n = malloc(4);
    keypair->e = malloc(4);
    keypair->d = malloc(4);
    
    if (!keypair->n || !keypair->e || !keypair->d) {
        rsa_free_keypair(keypair);
        return CRYPTO_ERROR_BUFFER_TOO_SMALL;
    }
    
    // 存储为大端序
    keypair->n[0] = (n >> 24) & 0xFF;
    keypair->n[1] = (n >> 16) & 0xFF;
    keypair->n[2] = (n >> 8) & 0xFF;
    keypair->n[3] = n & 0xFF;
    
    keypair->e[0] = (e >> 24) & 0xFF;
    keypair->e[1] = (e >> 16) & 0xFF;
    keypair->e[2] = (e >> 8) & 0xFF;
    keypair->e[3] = e & 0xFF;
    
    keypair->d[0] = (d >> 24) & 0xFF;
    keypair->d[1] = (d >> 16) & 0xFF;
    keypair->d[2] = (d >> 8) & 0xFF;
    keypair->d[3] = d & 0xFF;
    
    keypair->n_len = 4;
    keypair->e_len = 4;
    keypair->d_len = 4;
    
    return CRYPTO_SUCCESS;
}

int rsa_encrypt(const rsa_keypair_t *keypair, const uint8_t *plaintext, 
               size_t plaintext_len, uint8_t *ciphertext) {
    if (!keypair || !plaintext || !ciphertext) return CRYPTO_ERROR_INVALID_PARAM;
    
    // 重建n和e
    uint32_t n = ((uint32_t)keypair->n[0] << 24) | ((uint32_t)keypair->n[1] << 16) |
                 ((uint32_t)keypair->n[2] << 8) | keypair->n[3];
    uint32_t e = ((uint32_t)keypair->e[0] << 24) | ((uint32_t)keypair->e[1] << 16) |
                 ((uint32_t)keypair->e[2] << 8) | keypair->e[3];
    
    // 对每个字节进行加密
    for (size_t i = 0; i < plaintext_len && i < 4; i++) {
        uint32_t m = plaintext[i];
        uint32_t c = mod_pow(m, e, n);
        ciphertext[i] = c & 0xFF;
    }
    
    return CRYPTO_SUCCESS;
}

int rsa_decrypt(const rsa_keypair_t *keypair, const uint8_t *ciphertext,
               size_t ciphertext_len, uint8_t *plaintext) {
    if (!keypair || !ciphertext || !plaintext) return CRYPTO_ERROR_INVALID_PARAM;
    
    // 重建n和d
    uint32_t n = ((uint32_t)keypair->n[0] << 24) | ((uint32_t)keypair->n[1] << 16) |
                 ((uint32_t)keypair->n[2] << 8) | keypair->n[3];
    uint32_t d = ((uint32_t)keypair->d[0] << 24) | ((uint32_t)keypair->d[1] << 16) |
                 ((uint32_t)keypair->d[2] << 8) | keypair->d[3];
    
    // 对每个字节进行解密
    for (size_t i = 0; i < ciphertext_len && i < 4; i++) {
        uint32_t c = ciphertext[i];
        uint32_t m = mod_pow(c, d, n);
        plaintext[i] = m & 0xFF;
    }
    
    return CRYPTO_SUCCESS;
}

int rsa_sign(const rsa_keypair_t *keypair, const uint8_t *data, size_t data_len,
            uint8_t *signature) {
    if (!keypair || !data || !signature) return CRYPTO_ERROR_INVALID_PARAM;
    
    // 计算数据的哈希值
    uint8_t hash[SHA256_DIGEST_SIZE];
    int ret = sha256_hash(data, data_len, hash);
    if (ret != CRYPTO_SUCCESS) return ret;
    
    // 使用私钥对哈希值进行签名
    return rsa_decrypt(keypair, hash, SHA256_DIGEST_SIZE, signature);
}

int rsa_verify(const rsa_keypair_t *keypair, const uint8_t *data, size_t data_len,
              const uint8_t *signature) {
    if (!keypair || !data || !signature) return CRYPTO_ERROR_INVALID_PARAM;
    
    // 计算数据的哈希值
    uint8_t hash[SHA256_DIGEST_SIZE];
    int ret = sha256_hash(data, data_len, hash);
    if (ret != CRYPTO_SUCCESS) return ret;
    
    // 使用公钥验证签名
    uint8_t decrypted_hash[SHA256_DIGEST_SIZE];
    ret = rsa_encrypt(keypair, signature, SHA256_DIGEST_SIZE, decrypted_hash);
    if (ret != CRYPTO_SUCCESS) return ret;
    
    // 比较哈希值
    if (memcmp(hash, decrypted_hash, SHA256_DIGEST_SIZE) == 0) {
        return CRYPTO_SUCCESS;
    } else {
        return CRYPTO_ERROR_INVALID_SIGNATURE;
    }
}

void rsa_free_keypair(rsa_keypair_t *keypair) {
    if (keypair) {
        if (keypair->n) {
            crypto_secure_zero(keypair->n, keypair->n_len);
            free(keypair->n);
            keypair->n = NULL;
        }
        if (keypair->e) {
            crypto_secure_zero(keypair->e, keypair->e_len);
            free(keypair->e);
            keypair->e = NULL;
        }
        if (keypair->d) {
            crypto_secure_zero(keypair->d, keypair->d_len);
            free(keypair->d);
            keypair->d = NULL;
        }
        keypair->n_len = 0;
        keypair->e_len = 0;
        keypair->d_len = 0;
    }
} 